package com.lj.sse.utils;

import cn.hutool.core.collection.CollUtil;
import cn.hutool.json.JSONUtil;
import com.lj.common.exception.CommonException;
import com.lj.sse.result.SseRecord;
import lombok.extern.slf4j.Slf4j;
import org.springframework.http.MediaType;
import org.springframework.web.servlet.mvc.method.annotation.SseEmitter;

import java.io.IOException;
import java.util.Collection;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ScheduledFuture;

/**
 * @author luojing
 * @since 2025/3/26 14:16
 */
@Slf4j
public class SseContextUtils {

    /**
     * SseEmitter 的容器
     * key 为用户id
     * value 也是一个ConcurrentHashMap: key: 当前用户的当前客户端id value: SseRecord
     */
    private static final Map<Long, Map<String, SseRecord>> SSE_CONTEXT = new ConcurrentHashMap<>();


    public static SseRecord getSseRecord(Long userId, String clientId) {
        Map<String, SseRecord> sseRecordMap = SSE_CONTEXT.get(userId);
        if (CollUtil.isEmpty(sseRecordMap)) {
            return null;
        }
        return sseRecordMap.get(clientId);
    }

    public static void load(SseRecord sseRecord) {
        Map<String, SseRecord> sseRecordMap = SSE_CONTEXT.computeIfAbsent(sseRecord.getUserId(), userId -> new ConcurrentHashMap<>());
        if (sseRecordMap.get(sseRecord.getClientId()) != null) {
            throw new CommonException("连接已存在: " + sseRecord.getClientId());
        }
        sseRecordMap.put(sseRecord.getClientId(), sseRecord);
        // 为发射器注册相关回调事件
        SseEmitter emitter = sseRecord.getEmitter();
        // 关闭连接时回调
        emitter.onCompletion(() -> {
            log.info("用户({})连接({})关闭", sseRecord.getUserId(), sseRecord.getClientId());
            unload(sseRecord);
        });
        // 连接超时回调
        emitter.onTimeout(() -> {
            log.info("用户({})连接({})超时", sseRecord.getUserId(), sseRecord.getClientId());
            unload(sseRecord);
        });
        // 推送消息错误回调
        emitter.onError(e -> {
            log.info("用户({})连接({})出现错误: ", sseRecord.getUserId(), sseRecord.getClientId(), e);
            unload(sseRecord);
        });
        log.info("已成功装载用户({})的连接({})", sseRecord.getUserId(), sseRecord.getClientId());
    }

    public static void unload(SseRecord sseRecord) {
        unload(sseRecord.getUserId(), sseRecord.getClientId());
    }

    public static void unload(Long userId, String clientId) {
        Map<String, SseRecord> sseRecordMap = SSE_CONTEXT.get(userId);
        if (CollUtil.isEmpty(sseRecordMap) || sseRecordMap.get(clientId) == null) {
            // log.warn("未找到用户({})的连接({})进行卸载!", userId, clientId);
            return;
        }
        SseRecord sseRecord = sseRecordMap.remove(clientId);
        // 关闭定时
        ScheduledFuture<?> scheduledFuture = sseRecord.getScheduledFuture();
        if (scheduledFuture != null) {
            scheduledFuture.cancel(true);
        }
        log.info("已成功卸载用户({})的连接({})", userId, clientId);
    }

    /**
     * 向所有客户端发送消息
     *
     * @param message 消息
     */
    public static void sendAll(Object message) {
        Collection<Map<String, SseRecord>> values = SSE_CONTEXT.values();
        for (Map<String, SseRecord> sseRecordMap : values) {
            for (SseRecord sseRecord : sseRecordMap.values()) {
                send(sseRecord, message);
            }
        }
    }

    /**
     * 给指定用户的所有客户端发送消息
     *
     * @param userId  用户id
     * @param message 消息
     */
    public static void send(Long userId, Object message) {
        Map<String, SseRecord> sseRecordMap = SSE_CONTEXT.get(userId);
        if (CollUtil.isEmpty(sseRecordMap)) {
            log.warn("用户({})没有sse连接,无法发送消息:{}", userId, JSONUtil.toJsonStr(message));
            return;
        }
        for (SseRecord sseRecord : sseRecordMap.values()) {
            send(sseRecord, message);
        }
    }

    /**
     * 给指定用户的指定客户端发送消息
     *
     * @param userId   用户id
     * @param clientId 客户端id
     * @param message  消息
     */
    public static void send(Long userId, String clientId, Object message) {
        Map<String, SseRecord> sseRecordMap = SSE_CONTEXT.get(userId);
        if (CollUtil.isEmpty(sseRecordMap) || sseRecordMap.get(clientId) == null) {
            log.warn("未找到用户({})的连接({}) 无法发送消息:{}", userId, clientId, JSONUtil.toJsonStr(message));
            return;
        }
        SseRecord sseRecord = sseRecordMap.get(clientId);
        send(sseRecord, message);
    }

    /**
     * 根据SseRecord发送消息
     *
     * @param sseRecord sse记录
     * @param message   消息
     */
    public static void send(SseRecord sseRecord, Object message) {
        try {
            send(sseRecord.getEmitter(), message);
        } catch (Exception e) {
            log.error("向用户({})发送消息失败,断开连接({})", sseRecord.getUserId(), sseRecord.getClientId(), e);
            unload(sseRecord);
        }
    }

    /**
     * 发送消息给客户端
     *
     * @param sseEmitter sse发射器
     * @param message    消息，最终会以json的格式返回数据
     */
    private static void send(SseEmitter sseEmitter, Object message) throws IOException {
        SseEmitter.SseEventBuilder sseEventBuilder = SseEmitter.event().data(message, MediaType.APPLICATION_JSON);
        sseEmitter.send(sseEventBuilder);
    }
}
